#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
🎯 针对您系统的PyTorch CUDA安装脚本
根据您的CUDA 12.6版本提供安装指导
"""

print("🎯 针对您系统的PyTorch CUDA安装脚本")
print("=" * 80)

print("📋 您的系统信息:")
print("  GPU: NVIDIA GeForce RTX 3070")
print("  CUDA版本: 12.6")
print("  显存: 8GB")
print("  当前PyTorch版本: 2.8.0+cpu (仅CPU版本)")

print("\n🔧 推荐安装方案:")
print("  由于您的CUDA版本是12.6，推荐安装支持CUDA 12.1的PyTorch版本")

print("\n🚀 一键安装命令:")

print("\n1️⃣ 首先卸载当前的PyTorch版本:")
print("   pip uninstall torch torchvision torchaudio")

print("\n2️⃣ 安装支持CUDA 12.1的PyTorch版本:")
print("   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")

print("\n💡 为什么选择CUDA 12.1而不是12.6?")
print("   • PyTorch官方对CUDA 12.1支持最稳定")
print("   • CUDA 12.1向后兼容CUDA 12.6")
print("   • 性能和稳定性最佳")
print("   • 下载和安装速度更快")

print("\n⚡ 安装后性能提升:")
print("   • 排列5模型训练速度提升5-10倍")
print("   • 支持更大的batch_size (从16提升到32或更高)")
print("   • 可以训练更复杂的模型结构")
print("   • 减少训练时间从小时级到分钟级")

print("\n🛡️ 安装验证命令:")
print("   python -c \"import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA可用: {torch.cuda.is_available()}'); print(f'CUDA版本: {torch.version.cuda}' if torch.cuda.is_available() else '')\"")

print("\n🧪 安装后测试GPU训练:")
print("   1. 重新启动您的Python环境")
print("   2. 运行排列5训练脚本:")
print("      python scripts/plw/train_plw_sequence_lstm.py")
print("   3. 观察日志中是否显示使用CUDA设备")

print("\n🎯 安装成功标志:")
print("   日志中应该显示:")
print("   • 🎯 自动选择设备: cuda")
print("   • 🏷️ GPU名称: NVIDIA GeForce RTX 3070")
print("   • 📈 训练速度显著提升")

print("\n💡 小贴士:")
print("   • 如果pip安装慢，可以使用国内镜像源:")
print("     pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 -i https://pypi.tuna.tsinghua.edu.cn/simple/")
print("   • 安装完成后建议重启IDE或终端")
print("   • 如果遇到问题，可以先创建虚拟环境再安装")

print("\n" + "=" * 80)
print("🎉 安装完成后，您的排列5序列LSTM模型训练将自动使用GPU!")
print("🚀 训练时间将从几十分钟缩短到几分钟!")
print("🎯 享受GPU加速带来的极致体验!")